package com.spider.utils;

import com.spider.entites.Filed;
import org.apache.poi.hssf.usermodel.HSSFCell;
import org.apache.poi.hssf.usermodel.HSSFRow;
import org.apache.poi.hssf.usermodel.HSSFSheet;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * @author panglin
 * @date 2020/11/23 11:30
 */
public class ExcelExportUtils {

    public static ResponseEntity<byte[]> export(String taskName, List<Filed> filed, List<Map<String,String>> data){
        HSSFWorkbook workbook = new HSSFWorkbook();
        HSSFSheet sheet = workbook.createSheet("sheet1");
        for (int i = 0; i < filed.size(); i++) {
            sheet.setColumnWidth(i, 20*512);
        }
        HSSFRow row = sheet.createRow(0);
        row.setHeight((short) (25*20));
        for (int i = 0; i < filed.size(); i++) {
            HSSFCell cell = row.createCell(i);
            cell.setCellValue(filed.get(i).getFiled());
        }
        for (int i = 1; i < data.size() + 1; i++) {
            HSSFRow hssfRow = sheet.createRow(i);
            hssfRow.setHeight((short) (25*15));
            Map<String, String> map = data.get(i - 1);
            Set<String> keySet = map.keySet();
            int j = 0;
            for (String col : keySet) {
                HSSFCell cell = hssfRow.createCell(j);
                cell.setCellValue(map.get(col));
                j++;
            }
        }
        HttpHeaders httpHeaders = new HttpHeaders();
        try {
            httpHeaders.setContentDispositionFormData("attachment", new String((taskName + ".xls").getBytes(StandardCharsets.UTF_8), "GBK"));
        } catch (UnsupportedEncodingException e) {
            e.printStackTrace();
        }
        httpHeaders.setContentType(MediaType.APPLICATION_OCTET_STREAM);
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        try {
            workbook.write(bos);
            bos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return new ResponseEntity<>(bos.toByteArray(), httpHeaders, HttpStatus.CREATED);
    }
}
